#Analysis of Tajima's D, 26/08/10
#follows from BayesFst

#write function to perform sliding window analysis on a contig
#limit to snps with >100x coverage and >1% frequency

TajimaD <- function(data,sliding.window=1000,pos="start.pos",count="count",depth="depth",n=100){
	data <- data[data[,depth]>=n,c(pos,count,depth)]
	names(data) <- c('pos','count','depth')
	data$freq <- data$count/data$depth
	data <- data[data$freq>=0.01,]
	if(length(data$pos)!=length(unique(data$pos))){
		warning('same position shows more than one snp, only the first taken\n')
		data <- data[match(unique(data$pos),data$pos),]
		}
	D.out <- data.frame(pos=numeric(0),count=numeric(0),
		depth=numeric(0),freq=numeric(0),D=numeric(0),S=numeric(0))
	
	if(nrow(data)==0) return(D.out)
	data <- data[order(data$pos),]
	
	#left.edge <- data[1,'pos']
	#right.edge <- left.edge + window -1
	#last.snp <- data[nrow(data),'pos']
	
	#if(last.snp < right.edge) {
	#	warning('window too big for this contig\n')
	#	return(D.out) #bounces out with zero row dataframe
	#	}
	
	isnp <- 1
	for(isnp in 1:nrow(data)){
		
		left.edge <- data$pos[isnp] - sliding.window/2
		right.edge <- data$pos[isnp] + sliding.window/2
		
		tmp.data <- data[data$pos >=left.edge & data$pos <= right.edge, ]
		S <- nrow(tmp.data) #no. segregating sites
		dummy.count <- floor(n*tmp.data$freq)
		
		#check the following more carefully
		tmp <- (n-dummy.count)*dummy.count
		k.hat <- sum(tmp)/choose(n,2)
		a1 <- sum(1/seq(1,n-1))
		a2 <- sum(1/seq(1,n-1)^2)
		b1 <- (n+1)/(3*(n-1))
		b2 <- 2*(n^2 + n + 3)/(9*n*(n-1))
		c1 <- b1 - 1/a1
		c2 <- b2 - (n+2)/(a1*n)
		e1 <- c1/a1
		e2 <- c2/(a1^2 + a2)
		D <- (k.hat - S/a1)/sqrt(e1*S + e2*S*(S-1))
		D.out <- rbind(D.out,data.frame(
			pos=data[isnp,'pos'],count=data[isnp,'count'],
			depth=data[isnp,'depth'],freq=data[isnp,'freq'],D=D,S=S))
		
		#control for next window, remant of while loop
		#isnp <- isnp+1
		#left.edge <- data[isnp,'pos']
		#right.edge <- left.edge + window -1
		}
		
	D.out
	}

#now run over all genes
grouse.TjD.list <- vector('list',length=length(unique(grouse.df3$cDNA)))
names(grouse.TjD.list) <- unique(grouse.df3$cDNA)

for(gene in unique(grouse.df3$cDNA)){ #unique(grouse.df3$cDNA)
	tst.in <- grouse.df3[grouse.df3$cDNA==gene,]
	grouse.TjD.list[[gene]][[1]] <- try(
		TajimaD(tst.in,sliding.window=100,pos="start.pos",count="pop.1.count",depth="pop.1.tsv.depth",n=50)
		)
	grouse.TjD.list[[gene]][[2]] <- try(
		TajimaD(tst.in,sliding.window=100,pos="start.pos",count="pop.2.count",depth="pop.2.tsv.depth",n=50)
		)
	grouse.TjD.list[[gene]][[3]] <- try(
		TajimaD(tst.in,sliding.window=100,pos="start.pos",count="pop.3.count",depth="pop.3.tsv.depth",n=50)
		)
	names(grouse.TjD.list[[gene]]) <- c('pop1','pop2','pop3')
	}

#try to find interesting genes
#take D values from Tajima's paper
grouse.TjD.df <- data.frame(cDNA = unique(grouse.df3$cDNA),stringsAsFactors=FALSE)
grouse.TjD.df[,c('pop1.min','pop1.max','pop2.min','pop2.max','pop3.min','pop3.max')] <- 0

for(gene in 1:nrow(grouse.TjD.df)){
	if(!is.null(grouse.TjD.list[[gene]][[1]])&nrow(grouse.TjD.list[[gene]][[1]])>0) grouse.TjD.df[gene,c('pop1.min','pop1.max')] <- range(grouse.TjD.list[[gene]][[1]]$D)
	if(!is.null(grouse.TjD.list[[gene]][[2]])&nrow(grouse.TjD.list[[gene]][[2]])>0) grouse.TjD.df[gene,c('pop2.min','pop2.max')] <- range(grouse.TjD.list[[gene]][[2]]$D)
	if(!is.null(grouse.TjD.list[[gene]][[3]])&nrow(grouse.TjD.list[[gene]][[3]])>0) grouse.TjD.df[gene,c('pop3.min','pop3.max')] <- range(grouse.TjD.list[[gene]][[3]]$D)
	
	}